#include "model/root.h"
#include "macro.h"
#include "config.h"

#include <map>
#include <filesystem>

namespace iangels {

Ret ModelRoot::Init()
{
    IANGELS_CHECK(std::filesystem::exists(model_root_), FILE_NOT_EXIST,
        "model root: %s not exist", model_root_.c_str());

    IANGELS_CHECK_RET(config::load(model_root_ + "/config.json", config_),
        "load config.json failed");
    IANGELS_CHECK_RET(config::load(model_root_ + "/processor_config.json", processor_config_),
        "load processor_config.json failed");
    IANGELS_CHECK_RET(config::load(model_root_ + "/special_tokens_map.json", special_tokens_map_),
        "load special_tokens_map.json failed");
    IANGELS_CHECK_RET(config::load(model_root_ + "/tokenizer_config.json", tokenizer_config_),
        "load tokenizer_config.json failed");
    IANGELS_CHECK_RET(config::load(model_root_ + "/tokenizer.json", tokenizer_),
        "load tokenizer.json failed");

    return FindAllSafeTensorFiles();
}

Ret ModelRoot::FindAllSafeTensorFiles()
{
    for (const auto& entry : std::filesystem::directory_iterator(model_root_)) {
        if (entry.path().extension() == ".safetensors") {
            safetensor_files_.emplace_back(entry.path().string());
        }
    }

    IANGELS_CHECK(!safetensor_files_.empty(), SAFETENSOR_FILES_NOT_EXIST,
        "no safetensors file found in model root: %s", model_root_.c_str());

    return SUCCESS;
}

} // namespace iangels